import asyncio
import logging
from aiohttp import web
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, MediaStreamTrack
from aiortc.contrib.media import MediaBlackhole, MediaRecorder, MediaRelay
from go2_webrtc_driver.webrtc_driver import Go2WebRTCConnection, WebRTCConnectionMethod
from ultralytics import YOLO
import cv2
import numpy as np
from queue import Queue
import threading
import time
import json

frame_queue = Queue()

# WebRTC连接设置部分
conn = Go2WebRTCConnection(WebRTCConnectionMethod.LocalSTA, ip="192.168.123.161")

# Async function to receive video frames and put them in the queue
async def recv_camera_stream(track: MediaStreamTrack):
    while True:
        frame = await track.recv()
        # Convert the frame to a NumPy array
        img = frame.to_ndarray(format="bgr24")
        frame_queue.put(img)

def run_asyncio_loop(loop):
    asyncio.set_event_loop(loop)
    async def setup():
        try:
            # Connect to the device
            await conn.connect()

            # Switch video channel on and start receiving video frames
            conn.video.switchVideoChannel(True)

            # Add callback to handle received video frames
            conn.video.add_track_callback(recv_camera_stream)
        except Exception as e:
            logging.error(f"Error in WebRTC connection: {e}")

    # Run the setup coroutine and then start the event loop
    loop.run_until_complete(setup())
    loop.run_forever()

async def run():
    # Enable logging for debugging
    logging.basicConfig(level=logging.INFO)

    # 模型加载部分
    model = YOLO("yolo11n.pt")  # 替换为你的模型路径
    print("✅ YOLOv11n 模型加载成功，支持类别：", model.names)

    # Create a new event loop for the asyncio code
    loop = asyncio.new_event_loop()

    # Start the asyncio event loop in a separate thread
    asyncio_thread = threading.Thread(target=run_asyncio_loop, args=(loop,))
    asyncio_thread.start()
